﻿
#define SPADAS_DEPRECATED_HIDE
#include "NextGPUDecoder.h"
#include "ffmpegdl.h"

namespace videonext
{
	struct AVQSVFramesContext
	{
		Pointer surfaces;
		Int nb_surfaces;
		Int frame_type;
	};

	class GPUDecoderVars : public Vars
	{
	public:
		Pointer decoder;
		Int sentCount;
		Stream<Double> times;
		AVPixelFormat hardwareFormat;
		AVPixelFormat softwareFormat;
		Pointer hardwareDevice;
		Size2D frameSize;

		static AVPixelFormat getFormat(Pointer ctx, Array<AVPixelFormat> formats)
		{
			GPUDecoderVars *vars = (GPUDecoderVars*)avcodec_get_opaque(ctx);
			if (formats.contain(vars->hardwareFormat))
			{
				if (vars->hardwareFormat == AVPixelFormat::AV_PIX_FMT_QSV)
				{
					if (vars->frameSize.width == 0 || vars->frameSize.height == 0 || !initQSVFrames(ctx, vars->hardwareDevice, vars->frameSize))
					{
						if (formats.contain(AVPixelFormat::AV_PIX_FMT_NV12)) return AVPixelFormat::AV_PIX_FMT_NV12;
						else return AVPixelFormat::AV_PIX_FMT_NONE;
					}
				}
				return vars->hardwareFormat;
			}
			return AVPixelFormat::AV_PIX_FMT_NONE;
		}

	private:
		static Bool initQSVFrames(Pointer ctx, Pointer hwctx, Size2D frameSize)
		{
			#if FFMPEGDL_VER_MAJOR > 1 || FFMPEGDL_VER_MINOR >= 7
			Pointer hwFrameCtx = av_hwframe_ctx_alloc(hwctx);
			if (!hwFrameCtx) return FALSE;

			Pointer hwFrameCtxData = av_hwframe_ctx_get_data(hwFrameCtx);
			av_hwframe_ctx_data_set_format(hwFrameCtxData, AVPixelFormat::AV_PIX_FMT_QSV);
			av_hwframe_ctx_data_set_sw_format(hwFrameCtxData, avcodec_get_sw_pix_fmt(ctx));
			av_hwframe_ctx_data_set_width_align32(hwFrameCtxData, frameSize.width);
			av_hwframe_ctx_data_set_height_align32(hwFrameCtxData, frameSize.height);
			av_hwframe_ctx_data_set_initial_pool_size(hwFrameCtxData, 32);

			AVQSVFramesContext *qsvFramesCtx = (AVQSVFramesContext*)av_hwframe_ctx_data_get_hwctx(hwFrameCtxData);
			qsvFramesCtx->frame_type = 0x0010; // MFX_MEMTYPE_VIDEO_MEMORY_DECODER_TARGET

			Int ret = av_hwframe_ctx_init(hwFrameCtx);
			if (ret >= 0)
			{
				avcodec_set_hw_frames_ctx(ctx, hwFrameCtx);
				return TRUE;
			}
			else
			{
				av_buffer_unref(&hwFrameCtx);
				return FALSE;
			}
			#else
			return FALSE;
			#endif
		}
	};

	Array<Lock> hardwareCreateLock(12);
}

using namespace videonext;

GPUDecoder::GPUDecoder() : Object<GPUDecoderVars>(new GPUDecoderVars, TRUE)
{
	vars->decoder = 0;
	vars->hardwareDevice = 0;
}
bool GPUDecoder::open(VideoDataCodec codec, AVHWDeviceType type, Size2D size)
{
	close();

	String codecName;
	switch (codec)
	{
	case VideoDataCodec::MJPEG:
		codecName = "mjpeg";
		break;
	case VideoDataCodec::H264:
		codecName = "h264";
		break;
	case VideoDataCodec::H265:
		codecName = "hevc";
		break;
	default:
		return FALSE;
	}

	AVPixelFormat hardwareFormat = AVPixelFormat::AV_PIX_FMT_NONE;
	AVPixelFormat softwareFormat = AVPixelFormat::AV_PIX_FMT_NV12;
	UInt ffmpegLibVersion = getFFMPEGLibraryVersion();
	switch (type)
	{
	case AVHWDeviceType::AV_HWDEVICE_TYPE_CUDA:
		if (ffmpegLibVersion < 3) return FALSE;
		else if (ffmpegLibVersion < 4) codecName = codecName + "_cuvid";
		hardwareFormat = AVPixelFormat::AV_PIX_FMT_CUDA;
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_DXVA2:
		if (codec == VideoDataCodec::MJPEG) return FALSE;
		if (ffmpegLibVersion < 4) return FALSE;
		hardwareFormat = AVPixelFormat::AV_PIX_FMT_DXVA2_VLD;
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_D3D11VA:
		if (codec == VideoDataCodec::MJPEG) return FALSE;
		if (ffmpegLibVersion < 4) return FALSE;
		hardwareFormat = AVPixelFormat::AV_PIX_FMT_D3D11;
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_VAAPI:
		if (ffmpegLibVersion < 3) return FALSE;
		hardwareFormat = AVPixelFormat::AV_PIX_FMT_VAAPI;
		if (codec == VideoDataCodec::MJPEG) softwareFormat = AVPixelFormat::AV_PIX_FMT_YUV420P;
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_VDPAU:
		if (codec == VideoDataCodec::MJPEG) return FALSE;
		if (ffmpegLibVersion < 3) return FALSE;
		hardwareFormat = AVPixelFormat::AV_PIX_FMT_VDPAU;
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_VIDEOTOOLBOX:
		if (codec == VideoDataCodec::MJPEG) return FALSE;
		if (ffmpegLibVersion < 4) return FALSE;
		hardwareFormat = AVPixelFormat::AV_PIX_FMT_VIDEOTOOLBOX;
		break;
	case AVHWDeviceType::AV_HWDEVICE_TYPE_QSV:
		if (ffmpegLibVersion < 4) return FALSE;
		codecName = codecName + "_qsv";
		hardwareFormat = AVPixelFormat::AV_PIX_FMT_QSV;
		break;
	default:
		return FALSE;
	}

	Pointer targetCodec = avcodec_find_decoder_by_name(codecName.chars().data());
	if (targetCodec == 0) return false;

	Pointer decoderTmp = avcodec_alloc_context3(targetCodec);
	if (decoderTmp == 0) return false;

	Pointer hardwareTmp = NULL;
	hardwareCreateLock[(UInt)type].enter();
	Int resultCreateHardware = av_hwdevice_ctx_create(&hardwareTmp, type, NULL, NULL, 0);
	hardwareCreateLock[(UInt)type].leave();
	if (resultCreateHardware < 0)
	{
		avcodec_free_context(&decoderTmp);
		return false;
	}

	avcodec_set_hw_device_ctx(decoderTmp, hardwareTmp);
	avcodec_set_get_format(decoderTmp, vars->getFormat);
	avcodec_set_opaque(decoderTmp, vars);
	vars->hardwareFormat = hardwareFormat;

	Int resultOpenCodec = avcodec_open2(decoderTmp, targetCodec, NULL);
	if (resultOpenCodec < 0)
	{
		avcodec_free_context(&decoderTmp);
		return false;
	}

	vars->decoder = decoderTmp;
	vars->hardwareDevice = hardwareTmp;
	vars->sentCount = 0;
	vars->times = Stream<Double>(8);
	vars->frameSize = size;
	vars->softwareFormat = softwareFormat;

	return true;
}
void GPUDecoder::close()
{
	if (vars->decoder)
	{
		vars->hardwareDevice = 0;
		avcodec_free_context(&vars->decoder);
	}
}
Pointer GPUDecoder::getHardwareDevice()
{
	return vars->hardwareDevice;
}
Bool GPUDecoder::send(Binary packetData, Double time)
{
	if (!vars->decoder) return false;

	Pointer packet = av_packet_alloc();
	av_new_packet(packet, packetData.size());
	utility::memoryCopy(packetData.data(), av_packet_get_data(packet), packetData.size());

	av_packet_set_pts(packet, vars->sentCount);
	av_packet_set_dts(packet, vars->sentCount);

	Int resultSendPacket = avcodec_send_packet(vars->decoder, packet);
	Bool ok = resultSendPacket >= 0;

	if (ok)
	{
		vars->sentCount++;
		vars->times.enqueue(time);
	}

	av_packet_unref(packet);
	av_packet_free(&packet);

	return ok;
}
void GPUDecoder::flush()
{
	if (vars->decoder == 0) return;
	avcodec_send_packet(vars->decoder, NULL);
}
Bool GPUDecoder::receive(CompatibleImage& image, Double& time)
{
	if (!vars->decoder) return false;

	Pointer hardwareFrame = av_frame_alloc();
	if (!hardwareFrame) return false;

	Int resultReceiveFrame = avcodec_receive_frame(vars->decoder, hardwareFrame);
	if (resultReceiveFrame < 0)
	{
		if (resultReceiveFrame != av_get_error_code_value(AVErrorCode::AVERROR_EAGAIN) &&
			resultReceiveFrame != av_get_error_code_value(AVErrorCode::AVERROR_EOF)) // Hardware failure / 硬件失效
		{
			close();
		}
		av_frame_free(&hardwareFrame);
		return false;
	}
	else
	{
		if (vars->frameSize.width > 0 && av_frame_get_width(hardwareFrame) > (Int)vars->frameSize.width) av_frame_set_width(hardwareFrame, vars->frameSize.width);
		if (vars->frameSize.height > 0 && av_frame_get_height(hardwareFrame) > (Int)vars->frameSize.height) av_frame_set_height(hardwareFrame, vars->frameSize.height);

		if (av_frame_get_format(hardwareFrame) == vars->hardwareFormat)
		{
			Pointer hostFrame = av_frame_alloc();
			Int resultTransfer = av_hwframe_transfer_data(hostFrame, hardwareFrame, 0);
			if (resultTransfer >= 0 && av_frame_get_format(hostFrame) == vars->softwareFormat)
			{
				image = CompatibleImage(hostFrame, TRUE);
			}
			av_frame_free(&hardwareFrame);
		}
		else if (av_frame_get_format(hardwareFrame) == vars->softwareFormat)
		{
			image = CompatibleImage(hardwareFrame, TRUE);
		}
		else av_frame_free(&hardwareFrame);

		Array<Double> dequeued = vars->times.dequeue(1);
		Bool timeOK = dequeued.size() == 1;
		if (timeOK) time = dequeued[0];

		return image.isValid() && timeOK;
	}
}
Bool GPUDecoder::receiveHardwareFrame(Pointer& outputFrame, Double& time)
{
	if (!vars->decoder) return false;

	Pointer hardwareFrame = av_frame_alloc();
	if (!hardwareFrame) return false;

	Int resultReceiveFrame = avcodec_receive_frame(vars->decoder, hardwareFrame);
	if (resultReceiveFrame < 0)
	{
		if (resultReceiveFrame != av_get_error_code_value(AVErrorCode::AVERROR_EAGAIN) &&
			resultReceiveFrame != av_get_error_code_value(AVErrorCode::AVERROR_EOF)) // Hardware failure / 硬件失效
		{
			close();
		}
		av_frame_free(&hardwareFrame);
		return false;
	}
	else
	{
		if (vars->frameSize.width > 0 && av_frame_get_width(hardwareFrame) > (Int)vars->frameSize.width) av_frame_set_width(hardwareFrame, vars->frameSize.width);
		if (vars->frameSize.height > 0 && av_frame_get_height(hardwareFrame) > (Int)vars->frameSize.height) av_frame_set_height(hardwareFrame, vars->frameSize.height);

		Array<Double> dequeued = vars->times.dequeue(1);
		Bool timeOK = dequeued.size() == 1;

		if (timeOK && av_frame_get_format(hardwareFrame) == vars->hardwareFormat)
		{
			outputFrame = hardwareFrame;
			time = dequeued[0];
			return true;
		}

		av_frame_free(&hardwareFrame);
		return false;
	}
}

Bool GPUDecoder::downloadHardwareFrame(Pointer hardwareFrame, CompatibleImage& image, Bool releaseHardwareFrame)
{
	if (!vars->decoder) return false;
	if (!hardwareFrame) return false;
	if (av_frame_get_format(hardwareFrame) != vars->hardwareFormat) return false;

	Pointer hostFrame = av_frame_alloc();
	Int resultTransfer = av_hwframe_transfer_data(hostFrame, hardwareFrame, 0);

	Bool ok = resultTransfer >= 0 && av_frame_get_format(hostFrame) == vars->softwareFormat;
	if (ok) image = CompatibleImage(hostFrame, TRUE);

	if (releaseHardwareFrame) av_frame_free(&hardwareFrame);
	return ok;
}

Array<Lock> GPUDecoder::getHardwareCreateLocks()
{
	return hardwareCreateLock;
}